package com.toolbox.utils;

import org.opencv.core.*;
import org.opencv.dnn.Dnn;
import org.opencv.dnn.Net;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;
import org.opencv.imgproc.CLAHE;
import nu.pattern.OpenCV;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * 基于OpenCV的证件照背景更换工具
 * 使用深度学习模型实现精准人像分割和背景替换
 */
public class IdPhotoBackgroundChanger {

    private static final Logger logger = Logger.getLogger(IdPhotoBackgroundChanger.class.getName());

    static {
        // 使用OpenPnP的OpenCV自动加载机制
        nu.pattern.OpenCV.loadShared();
    }

    // 模型路径（实际使用时需要替换为有效的模型路径）
    private static final String MODEL_PATH = "models/deeplabv3.pb";
    private static final String MODEL_WEIGHTS_PATH = "models/deeplabv3.weights";
    
    // 标准证件照尺寸参数
    private static final int ID_PHOTO_WIDTH_1INCH = 295; // 1寸证件照宽度（像素）
    private static final int ID_PHOTO_HEIGHT_1INCH = 413; // 1寸证件照高度（像素）
    private static final int ID_PHOTO_WIDTH_2INCH = 413; // 2寸证件照宽度（像素）
    private static final int ID_PHOTO_HEIGHT_2INCH = 579; // 2寸证件照高度（像素）

    // 标准证件照红色背景（BGR格式）
    // 优化的证件照红色值，更加自然且符合标准
    private static final Scalar STANDARD_RED_BACKGROUND = new Scalar(50, 40, 210); // BGR格式 - 优化的证件照红色

    /**
     * 核心方法：将人物照片背景更改为红色
     * @param imageBytes 输入图像的字节数组
     * @param maxDimension 最大尺寸限制（像素）
     * @param enhanceFace 是否优化面部特征
     * @return 处理后的图像字节数组
     * @throws Exception 处理过程中的异常
     */
    public static byte[] changeToRedBackground(byte[] imageBytes, int maxDimension, boolean enhanceFace) throws Exception {
        // 1. 将字节数组转换为OpenCV的Mat对象
        Mat src = bytesToMat(imageBytes);
        if (src.empty()) {
            throw new Exception("无法加载图像数据");
        }

        try {
            // 2. 人像分割
            Mat mask = performHumanSegmentation(src);

            // 3. 创建红色背景
            Mat redBg = createRedBackground(src.size());

            // 4. 合成最终图像（保留前景，替换背景）
            Mat result = new Mat();
            src.copyTo(result, mask);
            // 使用位运算反转掩码，而不是矩阵求逆
            Mat invertedMask = new Mat();
            Core.bitwise_not(mask, invertedMask);
            redBg.copyTo(result, invertedMask);

            // 5. 优化图像质量和尺寸
            Mat optimized = optimizeIdPhoto(result, maxDimension, enhanceFace);

            // 6. 将Mat转换回字节数组
            return matToBytes(optimized);
        } catch (Exception e) {
            logger.log(Level.SEVERE, "处理图像失败", e);
            throw e;
        }
    }

    /**
     * 兼容性方法：支持旧版本API调用
     * @param imageBytes 输入图像的字节数组
     * @param targetBackground 目标背景色（在此实现中忽略，始终使用红色）
     * @return 处理后的图像字节数组
     * @throws Exception 处理过程中的异常
     */
    public static byte[] changePhotoBackground(byte[] imageBytes, String targetBackground) throws Exception {
        // 忽略targetBackground参数，始终使用红色背景
        // 使用默认的maxDimension=800和enhanceFace=true
        return changeToRedBackground(imageBytes, 800, true);
    }

    /**
     * 执行人像分割
     * @param src 原始图像
     * @return 人像掩码（白色为前景，黑色为背景）
     */
    private static Mat performHumanSegmentation(Mat src) {
        try {
            // 检查模型文件是否存在
            File modelFile = new File(MODEL_PATH);
            if (!modelFile.exists()) {
                logger.warning("深度学习模型文件不存在: " + MODEL_PATH);
                // 回退到传统的基于颜色的分割方法
                return fallbackSegmentation(src);
            }

            // 加载模型
            Net net = Dnn.readNetFromTensorflow(MODEL_PATH);
            
            // 预处理图像
            Mat resized = new Mat();
            Imgproc.resize(src, resized, new Size(513, 513));
            Mat blob = Dnn.blobFromImage(resized, 1.0, new Size(513, 513), new Scalar(127.5, 127.5, 127.5), true, false);
            
            // 设置输入并前向传播
            net.setInput(blob);
            Mat segMask = net.forward();
            
            // 后处理分割结果
            return processSegmentationMask(segMask, src.size());
        } catch (Exception e) {
            logger.log(Level.WARNING, "深度学习分割失败，使用传统方法", e);
            // 异常情况下使用传统方法
            return fallbackSegmentation(src);
        }
    }

    /**
     * 处理分割掩码
     * @param segMask 网络输出的分割掩码
     * @param originalSize 原始图像尺寸
     * @return 处理后的掩码
     */
    private static Mat processSegmentationMask(Mat segMask, Size originalSize) {
        // 重塑掩码
        Mat reshaped = segMask.reshape(1, (int) segMask.size(1));
        
        // 找到最大概率的类别
        Mat mask = new Mat(reshaped.size(0), reshaped.size(1), CvType.CV_8UC1);
        
        for (int i = 0; i < reshaped.rows(); i++) {
            for (int j = 0; j < reshaped.cols(); j++) {
                double[] data = reshaped.get(i, j);
                // COCO数据集中，人像类别ID为15
                mask.put(i, j, (data[15] > 0.5) ? 255 : 0);
            }
        }
        
        // 形态学操作，平滑边缘
        Mat kernel = Imgproc.getStructuringElement(Imgproc.MORPH_ELLIPSE, new Size(10, 10));
        Imgproc.morphologyEx(mask, mask, Imgproc.MORPH_CLOSE, kernel);
        
        // 调整回原始尺寸
        Imgproc.resize(mask, mask, originalSize, 0, 0, Imgproc.INTER_LINEAR);
        
        // 添加高斯模糊使边缘更平滑，避免锯齿状
        Imgproc.GaussianBlur(mask, mask, new Size(5, 5), 0);
        
        // 二值化处理，确保掩码清晰
        Imgproc.threshold(mask, mask, 127, 255, Imgproc.THRESH_BINARY);
        
        return mask;
    }

    /**
     * 备用分割方法：基于颜色和边缘检测
     * @param src 原始图像
     * @return 分割掩码
     */
    private static Mat fallbackSegmentation(Mat src) {
        // 1. 创建多个肤色检测通道
        // HSV色彩空间
        Mat hsv = new Mat();
        Imgproc.cvtColor(src, hsv, Imgproc.COLOR_BGR2HSV);
        
        // YCrCb色彩空间 - 对肤色检测更稳定
        Mat ycrcb = new Mat();
        Imgproc.cvtColor(src, ycrcb, Imgproc.COLOR_BGR2YCrCb);
        
        // 2. HSV肤色范围 - 扩大范围以适应更多肤色类型
        Scalar lowerSkinHSV = new Scalar(0, 15, 60);
        Scalar upperSkinHSV = new Scalar(25, 255, 255);
        Mat hsvMask = new Mat();
        Core.inRange(hsv, lowerSkinHSV, upperSkinHSV, hsvMask);
        
        // 3. YCrCb肤色范围
        Scalar lowerSkinYCrCb = new Scalar(0, 133, 77);
        Scalar upperSkinYCrCb = new Scalar(255, 173, 127);
        Mat ycrcbMask = new Mat();
        Core.inRange(ycrcb, lowerSkinYCrCb, upperSkinYCrCb, ycrcbMask);
        
        // 4. 合并两个掩码
        Mat skinMask = new Mat();
        Core.bitwise_or(hsvMask, ycrcbMask, skinMask);
        
        // 5. 形态学操作增强
        // 先膨胀再腐蚀，去除小洞
        Mat kernelClose = Imgproc.getStructuringElement(Imgproc.MORPH_ELLIPSE, new Size(10, 10));
        Imgproc.morphologyEx(skinMask, skinMask, Imgproc.MORPH_CLOSE, kernelClose);
        
        // 腐蚀再膨胀，去除小的噪点
        Mat kernelOpen = Imgproc.getStructuringElement(Imgproc.MORPH_ELLIPSE, new Size(5, 5));
        Imgproc.morphologyEx(skinMask, skinMask, Imgproc.MORPH_OPEN, kernelOpen);
        
        // 6. 额外的膨胀操作确保边缘完整
        Mat kernelDilate = Imgproc.getStructuringElement(Imgproc.MORPH_ELLIPSE, new Size(7, 7));
        Imgproc.dilate(skinMask, skinMask, kernelDilate);
        
        // 7. 边缘平滑 - 更大的高斯核和多次模糊
        Imgproc.GaussianBlur(skinMask, skinMask, new Size(9, 9), 0);
        Imgproc.GaussianBlur(skinMask, skinMask, new Size(7, 7), 0);
        
        // 8. 二值化处理，确保掩码清晰
        Imgproc.threshold(skinMask, skinMask, 100, 255, Imgproc.THRESH_BINARY);
        
        return skinMask;
    }

    /**
     * 创建红色背景
     * @param size 图像尺寸
     * @return 红色背景图像
     */
    private static Mat createRedBackground(Size size) {
        Mat redBg = new Mat(size, CvType.CV_8UC3, STANDARD_RED_BACKGROUND);
        // 添加轻微的纹理和渐变，使背景看起来更自然
        Mat gradient = new Mat(size, CvType.CV_8UC3);
        for (int i = 0; i < gradient.rows(); i++) {
            for (int j = 0; j < gradient.cols(); j++) {
                // 创建径向渐变效果
                double distance = Math.sqrt(Math.pow(i - gradient.rows()/2, 2) + Math.pow(j - gradient.cols()/2, 2));
                double maxDistance = Math.sqrt(Math.pow(gradient.rows()/2, 2) + Math.pow(gradient.cols()/2, 2));
                double intensity = Math.max(0, 1 - distance/maxDistance/2);
                double b = STANDARD_RED_BACKGROUND.val[0] * (1 + intensity * 0.2);
                double g = STANDARD_RED_BACKGROUND.val[1] * (1 + intensity * 0.2);
                double r = STANDARD_RED_BACKGROUND.val[2] * (1 + intensity * 0.2);
                gradient.put(i, j, Math.min(b, 255), Math.min(g, 255), Math.min(r, 255));
            }
        }
        
        // 添加非常细微的噪声
        Mat noise = Mat.zeros(size, CvType.CV_8UC1);
        Core.randn(noise, 0, 2); // 降低噪声强度
        Mat coloredNoise = new Mat();
        List<Mat> noiseChannels = new ArrayList<>();
        noiseChannels.add(noise);
        noiseChannels.add(noise);
        noiseChannels.add(noise);
        Core.merge(noiseChannels, coloredNoise);
        Core.add(gradient, coloredNoise, redBg);
        
        return redBg;
    }

    /**
     * 优化证件照质量和尺寸
     * @param image 输入图像
     * @param maxDimension 最大尺寸
     * @param enhanceFace 是否优化面部
     * @return 优化后的图像
     */
    private static Mat optimizeIdPhoto(Mat image, int maxDimension, boolean enhanceFace) {
        // 调整图像尺寸，保持比例
        Mat resized = resizePreservingAspectRatio(image, maxDimension);
        
        // 如果需要优化面部
        if (enhanceFace) {
            return enhanceImageQuality(resized);
        }
        
        return resized;
    }

    /**
     * 调整图像尺寸，保持宽高比
     * @param image 输入图像
     * @param maxDimension 最大尺寸
     * @return 调整后的图像
     */
    private static Mat resizePreservingAspectRatio(Mat image, int maxDimension) {
        double width = image.cols();
        double height = image.rows();
        double scale = Math.min(maxDimension / width, maxDimension / height);
        
        int newWidth = (int) (width * scale);
        int newHeight = (int) (height * scale);
        
        Mat resized = new Mat();
        // 使用高质量的插值方法，提升缩放后的图像清晰度
        Imgproc.resize(image, resized, new Size(newWidth, newHeight), 0, 0, Imgproc.INTER_CUBIC);
        
        return resized;
    }

    /**
     * 增强图像质量
     * @param image 输入图像
     * @return 增强后的图像
     */
    private static Mat enhanceImageQuality(Mat image) {
        Mat enhanced = new Mat();
        image.copyTo(enhanced);
        
        // 1. 应用双边滤波，在保留边缘的同时减少噪声
        Mat bilateralFiltered = new Mat();
        Imgproc.bilateralFilter(enhanced, bilateralFiltered, 9, 75, 75);
        
        // 2. 转换到Lab色彩空间，便于亮度和对比度调整
        Mat labImage = new Mat();
        Imgproc.cvtColor(bilateralFiltered, labImage, Imgproc.COLOR_BGR2Lab);
        
        // 分离通道
        List<Mat> labChannels = new ArrayList<>();
        Core.split(labImage, labChannels);
        
        // 3. 对L通道（亮度）进行自适应直方图均衡化，提升局部对比度
        CLAHE clahe = Imgproc.createCLAHE(2.0, new Size(8, 8));
        clahe.apply(labChannels.get(0), labChannels.get(0));
        
        // 合并通道
        Core.merge(labChannels, labImage);
        
        // 转换回BGR色彩空间
        Imgproc.cvtColor(labImage, enhanced, Imgproc.COLOR_Lab2BGR);
        
        // 4. 锐化图像 - 使用更强的锐化核
        Mat kernel = Mat.zeros(3, 3, CvType.CV_32F);
        kernel.put(1, 1, 6);  // 增强中心像素权重
        kernel.put(0, 1, -1);
        kernel.put(1, 0, -1);
        kernel.put(1, 2, -1);
        kernel.put(2, 1, -1);
        
        Imgproc.filter2D(enhanced, enhanced, -1, kernel);
        
        // 5. 轻度调整对比度和亮度
        Mat adjusted = new Mat();
        enhanced.convertTo(adjusted, -1, 1.1, 5); // alpha=1.1 (对比度), beta=5 (亮度)
        
        return adjusted;
    }

    /**
     * 将字节数组转换为OpenCV的Mat对象
     * @param imageBytes 图像字节数组
     * @return OpenCV Mat对象
     * @throws IOException 转换异常
     */
    private static Mat bytesToMat(byte[] imageBytes) throws IOException {
        return Imgcodecs.imdecode(new MatOfByte(imageBytes), Imgcodecs.IMREAD_COLOR);
    }

    /**
     * 将OpenCV的Mat对象转换为字节数组
     * @param mat OpenCV Mat对象
     * @return 图像字节数组
     */
    private static byte[] matToBytes(Mat mat) {
        MatOfByte buffer = new MatOfByte();
        Imgcodecs.imencode(".jpg", mat, buffer);
        return buffer.toArray();
    }

    /**
     * 生成标准1寸证件照
     * @param imageBytes 输入图像字节数组
     * @return 1寸证件照字节数组
     * @throws Exception 处理异常
     */
    public static byte[] generate1InchIdPhoto(byte[] imageBytes) throws Exception {
        Mat src = bytesToMat(imageBytes);
        Mat mask = performHumanSegmentation(src);
        
        // 创建1寸证件照尺寸的红色背景
        Mat redBg = createRedBackground(new Size(ID_PHOTO_WIDTH_1INCH, ID_PHOTO_HEIGHT_1INCH));
        
        // 调整人像大小，保持比例并居中
        Mat resizedPerson = new Mat();
        double aspectRatio = (double)src.cols() / src.rows();
        int newWidth, newHeight;
        if (aspectRatio > (double)ID_PHOTO_WIDTH_1INCH / ID_PHOTO_HEIGHT_1INCH) {
            newWidth = ID_PHOTO_WIDTH_1INCH;
            newHeight = (int)(ID_PHOTO_WIDTH_1INCH / aspectRatio);
        } else {
            newHeight = ID_PHOTO_HEIGHT_1INCH;
            newWidth = (int)(ID_PHOTO_HEIGHT_1INCH * aspectRatio);
        }
        Imgproc.resize(src, resizedPerson, new Size(newWidth, newHeight));
        
        // 调整掩码大小
        Mat resizedMask = new Mat();
        Imgproc.resize(mask, resizedMask, new Size(newWidth, newHeight));
        
        // 创建最终结果图像
        Mat finalResult = redBg.clone();
        
        // 计算居中位置
        int xOffset = (ID_PHOTO_WIDTH_1INCH - newWidth) / 2;
        int yOffset = (ID_PHOTO_HEIGHT_1INCH - newHeight) / 2;
        
        // 创建ROI（感兴趣区域）
        Rect roi = new Rect(xOffset, yOffset, newWidth, newHeight);
        
        // 将人像复制到红色背景上的居中位置
        resizedPerson.copyTo(finalResult.submat(roi), resizedMask);
        
        return matToBytes(enhanceImageQuality(finalResult));
    }

    /**
     * 生成标准2寸证件照
     * @param imageBytes 输入图像字节数组
     * @return 2寸证件照字节数组
     * @throws Exception 处理异常
     */
    public static byte[] generate2InchIdPhoto(byte[] imageBytes) throws Exception {
        Mat src = bytesToMat(imageBytes);
        Mat mask = performHumanSegmentation(src);
        
        // 创建2寸证件照尺寸的红色背景
        Mat redBg = createRedBackground(new Size(ID_PHOTO_WIDTH_2INCH, ID_PHOTO_HEIGHT_2INCH));
        
        // 调整人像大小，保持比例并居中
        Mat resizedPerson = new Mat();
        double aspectRatio = (double)src.cols() / src.rows();
        int newWidth, newHeight;
        if (aspectRatio > (double)ID_PHOTO_WIDTH_2INCH / ID_PHOTO_HEIGHT_2INCH) {
            newWidth = ID_PHOTO_WIDTH_2INCH;
            newHeight = (int)(ID_PHOTO_WIDTH_2INCH / aspectRatio);
        } else {
            newHeight = ID_PHOTO_HEIGHT_2INCH;
            newWidth = (int)(ID_PHOTO_HEIGHT_2INCH * aspectRatio);
        }
        Imgproc.resize(src, resizedPerson, new Size(newWidth, newHeight));
        
        // 调整掩码大小
        Mat resizedMask = new Mat();
        Imgproc.resize(mask, resizedMask, new Size(newWidth, newHeight));
        
        // 创建最终结果图像
        Mat finalResult = redBg.clone();
        
        // 计算居中位置
        int xOffset = (ID_PHOTO_WIDTH_2INCH - newWidth) / 2;
        int yOffset = (ID_PHOTO_HEIGHT_2INCH - newHeight) / 2;
        
        // 创建ROI（感兴趣区域）
        Rect roi = new Rect(xOffset, yOffset, newWidth, newHeight);
        
        // 将人像复制到红色背景上的居中位置
        resizedPerson.copyTo(finalResult.submat(roi), resizedMask);
        
        return matToBytes(enhanceImageQuality(finalResult));
    }
}